package MLDA_Process;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.apache.commons.math3.special.Beta;
import org.apache.commons.math3.special.Gamma;


import Commons.IOUtil;
import gnu.trove.set.hash.TIntHashSet;

public class DualSparseTM implements Runnable {

    public int K = 100;

    public int M;
    public int V;

    public double alpha0 = 1E-7;
    public double alpha1 = 0.1;
    public double beta0 = 1E-7;
    public double beta1 = 0.01;

    public double gamma0 = 1.0;
    public double gamma1 = 1.0;

    public int niters = 2000;

    public int nwk[][];
    public int nk[];

    public int ndk[][];
    public int nd[];

    public TIntHashSet ndkSet[];
    public TIntHashSet nkwSet[];

    public int z_assigns[][];

    public String inputPath;
    public String outputPath;

    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();

    public void loadTxts(String txtPath) {
        BufferedReader reader = IOUtil.getReader(txtPath, "UTF-8");

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();

                String[] tokens = line.trim().split("\\s+");
                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        M = docs.size();
        V = w2i.size();

        return;
    }

    public void init_model() {
        nwk = new int[V][K];
        nk = new int[K];

        ndk = new int[M][K];
        nd = new int[M];

        z_assigns = new int[M][];

        ndkSet = new TIntHashSet[M];
        nkwSet = new TIntHashSet[K];

        for (int k = 0; k != K; k++) {
            nkwSet[k] = new TIntHashSet();
        }

        for (int m = 0; m != M; m++) {
            int N = docs.get(m).size();
            z_assigns[m] = new int[N];

            ndkSet[m] = new TIntHashSet();
            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                int z = (int) Math.floor(Math.random() * K);

                ndkSet[m].add(z);

                nwk[w][z]++;
                nk[z]++;
                ndk[m][z]++;
                nd[m]++;

                nkwSet[z].add(w);

                z_assigns[m][n] = z;
            }
        }
    }

    public void sampleTopic(int m, int n) {
        int w = docs.get(m).get(n);
        int z = z_assigns[m][n];

        int S = ndkSet[m].size();
        int R = nkwSet[z].size();

        ndk[m][z]--;
        nd[m]--;
        nwk[w][z]--;
        nk[z]--;

        double[] pTable = new double[K];
        int x = 0, y = 0;
        for (int k = 0; k != K; k++) {
            if (ndkSet[m].contains(k)) x = 1;
            else x = 0;

            if (nkwSet[k].contains(w)) y = 1;
            else y = 0;

            pTable[k] = (ndk[m][k] + x * alpha1 + alpha0) / (nd[m] + K * alpha0 + S * alpha1)
                    * (nwk[w][k] + y * beta1 + beta0) / (nk[k] + V * beta0 + R * beta1);
        }

        for (int k = 1; k != K; k++) {
            pTable[k] += pTable[k - 1];
        }

        double r = Math.random() * pTable[K - 1];

        for (int k = 0; k != K; k++) {
            if (r < pTable[k]) {
                z = k;
                break;
            }
        }

        ndk[m][z]++;
        nd[m]++;
        nwk[w][z]++;
        nk[z]++;

        z_assigns[m][n] = z;
    }

    public void sampleTopicSelectors(int m) {

        int N = docs.get(m).size();

        for (int k = 0; k != K; k++) {
            if (ndkSet[m].contains(k)) {
                ndkSet[m].remove(k);
            }

            int S = ndkSet[m].size();

            double[] pTable = new double[2];

            pTable[0] = (K - S - 1.0 + gamma0) * Gamma.gamma(alpha1 + alpha0)
                    * Math.exp(Beta.logBeta(K * gamma0 + S * gamma1, N + S * gamma1 + gamma1 + K * gamma0));

            pTable[1] = (S + gamma1) * Gamma.gamma(ndk[m][k] + alpha0 + alpha1)
                    * Math.exp(Beta.logBeta(K * gamma0 + S * gamma1 + gamma1, N + S * gamma1 + K * gamma0));

            pTable[1] += pTable[0];
            double r = Math.random() * pTable[1];

            if (r > pTable[0]) {
                ndkSet[m].add(k);
            }

        }
    }


    public void sampleTermSelectors(int k) {
        int N = nk[k];

        for (int w = 0; w != V; w++) {
            if (nkwSet[k].contains(w)) {
                nkwSet[k].remove(w);
            }

            int R = nkwSet[k].size();

            double[] pTable = new double[2];

            pTable[0] = (V - R - 1.0 + gamma0) * Gamma.gamma(beta1 + beta0)
                    * Math.exp(Beta.logBeta(V * gamma0 + R * gamma1, N + R * gamma1 + gamma1 + V * gamma0));

            pTable[1] = (R + gamma1) * Gamma.gamma(nwk[w][k] + beta0 + beta1)
                    * Math.exp(Beta.logBeta(V * gamma0 + R * gamma1 + gamma1, N + R * gamma1 + V * gamma0));

            pTable[1] += pTable[0];
            double r = Math.random() * pTable[1];

            if (r > pTable[0]) {
                nkwSet[k].add(w);
            }
        }
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("Iteration " + iter + " ...");

            Long begin = System.currentTimeMillis();
            for (int m = 0; m != M; m++) {
                for (int n = 0; n != docs.get(m).size(); n++) {
                    sampleTopic(m, n);
                }
            }

            for (int m = 0; m != M; m++) {
                sampleTopicSelectors(m);
            }

            for (int k = 0; k != K; k++) {
                sampleTermSelectors(k);
            }

            System.out.println((System.currentTimeMillis() - begin) / 1000.0);
        }

        double sum = 0.0;
        for (int m = 0; m != M; m++) {
            sum += ndkSet[m].size();
        }
        System.out.println(sum / (double) M);
    }

    public double[][] computePhi() {
        double[][] phi = new double[K][V];
        int y = 0;
        int R = 0;
        for (int k = 0; k != K; k++) {
            R = nkwSet[k].size();
            for (int v = 0; v != V; v++) {
                if (nkwSet[k].contains(v)) y = 1;
                else y = 0;
                phi[k][v] = (nwk[v][k] + y * beta1 + beta0) / (nk[k] + V * beta0 + R * beta1);
            }
        }
        return phi;
    }

    public double[][] computeTheta() {
        double[][] theta = new double[M][K];
        int x = 0;
        int S = 0;
        for (int m = 0; m != M; m++) {
            S = ndkSet[m].size();
            for (int k = 0; k != K; k++) {
                if (ndkSet[m].contains(k)) x = 1;
                else x = 0;
                theta[m][k] = (ndk[m][k] + x * alpha1 + alpha0) / (nd[m] + S * alpha1 + K * alpha0);
            }
        }
        return theta;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });

            res.add(pairs);
        }

        return res;
    }

    public void printTopics(String path, int top_n) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(path);
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            writer.write("Topic " + k + ":\n");
            for (int i = 0; i != top_n; i++) {
                writer.write(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue() + "\n");
            }
        }
        writer.close();
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] phi = computePhi();

        try {
            for (int k = 0; k != K; k++) {
                for (int v = 0; v != V; v++) {
                    writer.append(phi[k][v] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveTheta(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] theta = computeTheta();

        try {
            for (int m = 0; m != M; m++) {
                for (int k = 0; k != K; k++) {
                    writer.append(theta[m][k] + " ");
                }
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveZAssigns(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (int m = 0; m != M; m++) {
                for (int n = 0; n != docs.get(m).size(); n++) {
                    writer.append(docs.get(m).get(n) + ":" + z_assigns[m][n] + " ");
                }
                writer.newLine();
            }

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

//	public static void main(String args[]) {
//		String root= "D:\\javaEE\\MLDA\\data2\\news\\";
//		DualSparseTM stm = new DualSparseTM();
//		stm.loadTxts(root+"en09062011.news.data.processed");///Users/zuoyuan/Desktop/Experiment/20130601-origin
//		stm.init_model();
//		stm.estimate();
//		stm.saveTheta(root+"model-3.theta");
//		stm.printTopics(10);
//	}

    public void run() {
        this.loadTxts(this.inputPath);///Users/zuoyuan/Desktop/Experiment/20130601-origin
        this.init_model();
        this.estimate();
        try {
            this.printTopics(outputPath + "/model-final.twords", 20);
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }
        this.saveWordmap(outputPath + "/wordmap.txt");
        this.savePhi(outputPath + "/model-final.phi");
        this.saveZAssigns(outputPath + "/model-final.tassign");
        this.saveTheta(outputPath + "/model-final.theta");
    }

    public DualSparseTM(int iter, String inputPath, String outputPath) {
        this.niters = iter;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }

    public static void DualSparseTMs(int iter, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath() + "/DualSparseTM/";
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + i)).mkdirs();
            service.submit(new DualSparseTM(iter, path, parent_path + i));

        }
        service.shutdown();
    }
}
